from fastapi import APIRouter, HTTPException, UploadFile, File
from app.schemas.dataset import DatasetUpload, DatasetInfo, TrainingConfig, TrainingResponse
from app.core.config import settings
import pandas as pd
import json
from pathlib import Path

router = APIRouter()

@router.post("/upload", response_model=DatasetInfo)
async def upload_dataset(dataset: DatasetUpload):
    """上传数据集"""
    try:
        # 创建DataFrame
        df = pd.DataFrame([
            {"text": email.text, "label": email.label}
            for email in dataset.emails
        ])
        
        # 保存数据集
        output_path = settings.DATA_DIR / f"{dataset.dataset_name}.csv"
        df.to_csv(output_path, index=False)
        
        return DatasetInfo(
            dataset_name=dataset.dataset_name,
            total_count=len(df),
            spam_count=len(df[df['label'] == 'spam']),
            ham_count=len(df[df['label'] == 'ham'])
        )
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.post("/upload/file")
async def upload_dataset_file(
    dataset_name: str,
    file: UploadFile = File(...)
):
    """上传CSV文件作为数据集"""
    try:
        # 读取上传的CSV文件
        content = await file.read()
        df = pd.read_csv(pd.io.common.BytesIO(content))
        
        # 验证必要的列
        if not all(col in df.columns for col in ['text', 'label']):
            raise ValueError("CSV必须包含'text'和'label'列")
        
        # 保存数据集
        output_path = settings.DATA_DIR / f"{dataset_name}.csv"
        df.to_csv(output_path, index=False)
        
        return DatasetInfo(
            dataset_name=dataset_name,
            total_count=len(df),
            spam_count=len(df[df['label'] == 'spam']),
            ham_count=len(df[df['label'] == 'ham'])
        )
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.get("/list")
async def list_datasets():
    """列出所有可用的数据集"""
    try:
        datasets = []
        for file in settings.DATA_DIR.glob("*.csv"):
            df = pd.read_csv(file)
            datasets.append(
                DatasetInfo(
                    dataset_name=file.stem,
                    total_count=len(df),
                    spam_count=len(df[df['label'] == 'spam']),
                    ham_count=len(df[df['label'] == 'ham'])
                )
            )
        return datasets
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e)) 